// Copyright (c) Lawrence Livermore National Security, LLC and other VisIt
// Project developers.  See the top-level LICENSE file for dates and other
// details.  No copyright assignment is required to contribute to VisIt.
// ************************************************************************* //
//  File: avtFFTFilter.C
// ************************************************************************* //

#include <avtFFTFilter.h>
#include <vtkDataSet.h>
#include <vtkFloatArray.h>
#include <vtkDoubleArray.h>
#include <vtkCell.h>
#include <vtkVisItUtility.h>
#include <vtkCellData.h>
#include <vtkPointData.h>
#include <vtkRectilinearGrid.h>
#include <vtkPoints.h>
#include <iostream>
#include <fstream>
#include <DebugStream.h>

#include <algorithm>
#include <math.h>

#if defined(_MSC_VER)
#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif
#endif

#define SWAP(a, b) tempr=(a); (a)=(b); (b)=tempr


// ****************************************************************************
//  Method: avtFFTFilter constructor
//
//  Programmer: bala -- generated by xml2avt
//  Creation:   Thu Aug 18 09:28:06 PDT 2011
//
// ****************************************************************************

avtFFTFilter::avtFFTFilter()
{
}


// ****************************************************************************
//  Method: avtFFTFilter destructor
//
//  Programmer: bala -- generated by xml2avt
//  Creation:   Thu Aug 18 09:28:06 PDT 2011
//
//  Modifications:
//
// ****************************************************************************

avtFFTFilter::~avtFFTFilter()
{
}


// ****************************************************************************
//  Method:  avtFFTFilter::Create
//
//  Programmer: bala -- generated by xml2avt
//  Creation:   Thu Aug 18 09:28:06 PDT 2011
//
// ****************************************************************************

avtFilter *
avtFFTFilter::Create()
{
    return new avtFFTFilter();
}


// ****************************************************************************
//  Method:      avtFFTFilter::SetAtts
//
//  Purpose:
//      Sets the state of the filter based on the attribute object.
//
//  Arguments:
//      a        The attributes to use.
//
//  Programmer: bala -- generated by xml2avt
//  Creation:   Thu Aug 18 09:28:06 PDT 2011
//
// ****************************************************************************

void
avtFFTFilter::SetAtts(const AttributeGroup *a)
{
    atts = *(const FFTAttributes*)a;
}


// ****************************************************************************
//  Method: avtFFTFilter::Equivalent
//
//  Purpose:
//      Returns true if creating a new avtFFTFilter with the given
//      parameters would result in an equivalent avtFFTFilter.
//
//  Programmer: bala -- generated by xml2avt
//  Creation:   Thu Aug 18 09:28:06 PDT 2011
//
// ****************************************************************************

bool
avtFFTFilter::Equivalent(const AttributeGroup *a)
{
    return (atts == *(FFTAttributes*)a);
}


// ****************************************************************************
//  Method: avtFFTFilter::Fourier
//
//  Purpose:
//      Performs FFT operation.
//      Replace `data' by its discrete Fourier transform, if `isign'is
//      input as 1, or by its inverse discrete Fourier transform, if 
//      `isign' is input as -1. `data' is a complex array of length `nn',
//      input as a real array data[1..2*nn]. `nn' MUST be an integer
//      power of 2 (this is not checked for!?)  */
//
//  Programmer: bala 
//  Creation:   Thu Aug 18 09:28:06 PDT 2011
//
//  Changed by mdurant Aug 24 2011
//    This code was written in fortran-style
//    i.e. with array indices from [1 ... N]
//    Rather than monkey with the actual values of the loop control variables
//    I have solved this by subtracting 1 every time we access the "data" array
// ****************************************************************************


void avtFFTFilter::Fourier(double* data, long nn, long isign)
{
  int i, j, m, n;
  int  mmax, istep;
  float wtemp, wr, wpr, wpi, wi, theta;
  float tempr, tempi;
  
  long dataLength = nn * 2;
  debug5 <<"avtFFTFilter::Fourier() - DataLength is " <<dataLength <<std::endl;
  n= nn<<1;
  j=1;
  
  for(i=1; i<n; i+=2)
  {
    if(j > i)
    {
      SWAP(data[j - 1], data[i - 1]);
      SWAP(data[j+1 - 1], data[i+1 - 1]);
    }
    m= n>>1;
    while(m >=2 &&j >m)
    {
      j -=m;
      m >>=1;
    }
    j +=m;
  }
  mmax =2;
  while(n> mmax)
  {
    istep= 2*mmax;
    theta= 6.28318530717959/(isign*mmax);
    wtemp= sin(.5*theta);
    wpr= -2.0*wtemp*wtemp;
    wpi= sin(theta);
    wr= 1.0;
    wi= 0.0;
    for(m=1; m<mmax; m+=2)
    {
      for(i=m; i<n; i+=istep)
      {
        j=i+mmax;

        if ((j + 1 - 1) >= dataLength) {
          debug5 <<"avtFFTFilter::Fourier() - ERROR: j + 1 - 1 exceeded dataLength in Fourier function." <<std::endl;
          continue;
        }

        tempr= wr*data[j - 1]-wi*data[j+1 - 1];
        tempi= wr*data[j+1 - 1]+wi*data[j - 1];
        data[j - 1]= data[i - 1]- tempr;
        data[j+1 - 1]= data[i+1 - 1]- tempi;
        data[i - 1] +=tempr;
        data[i+1 - 1] +=tempi;
      }
      wr= (wtemp=wr)*wpr - wi*wpi+wr;
      wi= wi*wpr + wtemp*wpi + wi;
    }
    mmax= istep;
  }
}


// **************************************************************************
//  Method: avtFFTFilter::PopulateRGrid
//
//  Purpose:
//      Prepares the data for plotting. Populates the Rectilinear Grid with 
//  frequency data for the X axis and fftData for the Y axis 
//
//  Arguments:
//      freqs      Frequency data.
//      fftData      The computed FFT data.
//      nFFThalf    The number of points for both X and Y axis
//      rGrid      The Rectilinear Grid that will be populated with the 
//        frequency/FFT data.
//
//  Programmer: bala 
//  Creation:   Mon Aug 16 13:48:35 PST 2011
// **************************************************************************

vtkRectilinearGrid* avtFFTFilter::PopulateRGrid(std::vector<double>* freqs, double* fftData, long nFFThalf) {
    debug5 <<"avtFFTFilter::PopulateRGrid() - Entering." <<std::endl;
  
    int nFFTHalf = (int)nFFThalf; 
    
    // Create an output dataset
    vtkFloatArray* vals = vtkFloatArray::New();
    vals->SetNumberOfComponents(1);
    vals->SetNumberOfTuples(nFFTHalf);

    vtkRectilinearGrid* outGrid = vtkVisItUtility::Create1DRGrid(nFFTHalf, VTK_FLOAT);
    outGrid->GetPointData()->SetScalars(vals);
    vtkFloatArray* outGrid_X = vtkFloatArray::SafeDownCast(outGrid->GetXCoordinates());
    outGrid_X->SetNumberOfTuples(nFFTHalf);

    double* dummy = new double[1];

    for (int i = 0 ; i < nFFTHalf; i++) {
      outGrid_X->SetTuple1(i, freqs->operator[](i));
      *dummy = fftData[i];
      vals->SetTuple(i, dummy);
    }

    delete []dummy;
    vals->Delete();

    debug5 <<"avtFFTFilter::PopulateRGrid() - Returning" <<std::endl;
    return outGrid;
}


// **************************************************************************
//  Method: avtFFTFilter::ComputeFFT
//
//  Purpose:
//      Operates the FFT filter on the input data.
//
//  Arguments:
//  data  The input data 
//  dt  The value of the dt
//  rGrid  Pointer to the output dataset.
//
//  Programmer: bala 
//  Creation:   Mon Aug 16 13:48:35 PST 2011
// 
//  Psuedo code that computes the FFT:
//
//  Given (history) data[0:nData-1], and dt.
// 
//  1) find power of 2 greater than length of history data:  nPwr2 = Ceiling(log(nData)/log(2.0))
//  2) Get length of work array:  nFFT = 2^nPwr2, nFFTdouble = 2*nFFT, nFFThalf = nFFT/2
//  3) new double fftData[nFFTdouble], and zero-fill.
//  4) copy data over to fftData, with stride 2
//     fftData[0] = data[0]
//     fftData[2] = data[1]
//     fftData[4] = data[2]
//     ... repeat nData times.
// 
//  5) perform the FFT, four1(fftData,nFFT,1)
//  6) compute complex amplitude, dB, and compress in array
//     fftData[0] = 10*log10(max(1.0e-25,fftData[0]^2+fftData[1]^2))
//     fftData[1] = 10*log10(max(1.0e-25,fftData[2]^2+fftData[3]^2))
//     fftData[2] = 10*log10(max(1.0e-25,fftData[4]^2+fftData[5]^2))
//     ... repeat nFFThalf times.
// 
//  7) df = 1.0/(nFFT*dt)
//  8) get frequency axis:  new freqs[nFFThalf], freqs[i]=i*df
//  9) Plot fftData[0:nFFThalf-1] vs. freqs.
// **************************************************************************

vtkRectilinearGrid* avtFFTFilter::ComputeFFT(std::vector<double>* data, long nData, double dt) {
  debug5 <<"avtFFTFilter::ComputeFFT() - Entering" <<std::endl;

  // 1. Find the power of 2 greater than length of history data
  //We make a variable containing the value 2.0
  //So that the windows compiler knows which "log" method we're trying to call
  double doubleTwo = 2.0;
  long nPwr2 = ceil(log((double)nData)/log(doubleTwo));

  // 2) Get length of work array:  nFFT = 2^nPwr2, nFFTdouble = 2*nFFT, nFFThalf = nFFT/2
  long longTwo = 2;
  long nFFT = pow((long double)longTwo, (int)nPwr2);
  long nFFTdouble = 2 * nFFT;
  long nFFThalf = nFFT / 2;

  // 3) new double fftData[nFFTdouble], and zero-fill.
  double* fftData = new double[nFFTdouble];
  for(long i = 0; i < nFFTdouble; i++) {
    fftData[i] = 0;
  }

  // 4) copy data over to fftData, with stride 2
  for(long i = 0; i < nData; i++) {
    fftData[i * 2] = data->operator[](i);
  }

  // 5) perform the FFT, Fourier(fftData,nFFT,1)
  Fourier(fftData, nFFT, 1);

  // 6) compute complex amplitude, dB, and compress in array
  double oneENeg25 = pow((double)10, -25);

  for(int i = 0; i < nFFThalf; i++) {
    int jIndex = i * 2;
    fftData[i] = 10 * log10(std::max(oneENeg25, (pow(fftData[jIndex], 2) + pow(fftData[jIndex + 1], 2))));
  }

  // 7) df = 1.0/(nFFT*dt)
  double df = 1.0/(nFFT * dt);

  // 8) get frequency axis:  new freqs[nFFThalf], freq[i]=i*df
  std::vector<double> freqs;
  freqs.reserve(nFFThalf);
  for(long i = 0; i < nFFThalf; i++) {
    freqs.push_back(i * df);
  }  
    
  // 9) Plot fftData[0:nFFThalf-1] vs. freqs.
  vtkRectilinearGrid* outGrid = PopulateRGrid(&freqs, fftData, nFFThalf);

  delete []fftData;

  debug5 <<"avtFFTFilter::ComputeFFT() - Returning" <<std::endl;
  return outGrid;
}


// ****************************************************************************
//  Method: avtFFTFilter::ExecuteData
//
//  Purpose:
//      Sends the specified input and output through the FFT filter.
//
//  Arguments:
//      in_dr      The input data representation.
//
//  Returns:       The output data representation.
//
//  Programmer: bala
//  Creation:   Mon Aug 15 13:48:35 PST 2011
//
//  Modifications:
//    Eric Brugger, Mon Jul 28 15:22:58 PDT 2014
//    Modified the class to work with avtDataRepresentation.
//
// ****************************************************************************

avtDataRepresentation *
avtFFTFilter::ExecuteData(avtDataRepresentation *in_dr)
{
  debug5 <<"avtFFTFilter::ExecuteData() - Entering." <<std::endl;

  //
  // Get the VTK data set.
  //
  vtkDataSet *in_ds = in_dr->GetDataVTK();

  // Check if the input dataset is a rectilinear grid.
  if (in_ds->GetDataObjectType() != VTK_RECTILINEAR_GRID) {
    debug3 << "avtFFTFilter::ExecuteData() ERROR: Input dataset is not of type VTK_RECTILINEAR_GRID." <<std::endl;
    debug3 << "avtFFTFilter::ExecuteData() ERROR: Not executing FFT. Returning input dataset." << std::endl;
    return in_dr;
  }

  // Convert input dataset to a rectilinear grid for easier access
  vtkRectilinearGrid *rGrid = vtkRectilinearGrid::SafeDownCast(in_ds);

  // If input dataset has not been safely down cast to vtkRectilinearGrid, return.
  if(rGrid == NULL) {
    debug3 << "avtFFTFilter::ExecuteData() ERROR: Input dataset when safely down cast to VTK_RECTILINEAR_GRID is NULL." <<std::endl;
    debug3 << "avtFFTFilter::ExecuteData() ERROR: Not executing FFT. Returning input dataset." << std::endl;
    return in_dr;
  }

  // Check if the input dataset dimension is X/1/1
  int* dims = rGrid->GetDimensions();
  if ((dims[1] != 1) || (dims[2] != 1)) {
    debug3 << "avtFFTFilter::ExecuteData() ERROR: The dimension of the input dataset is not 1 (i.e., [x][1][1])." <<std::endl;
    debug3 << "avtFFTFilter::ExecuteData() ERROR: Not executing FFT. Returning input dataset." << std::endl;
    return in_dr;
  }

  // Get the scalars.
  vtkDataArray* scalars = rGrid->GetPointData()->GetScalars();
  unsigned int nPts = scalars->GetNumberOfTuples();

  // Not operating if the no of tuples in the scalars is <= 1.
  if (nPts <= 1) {
    debug3 << "avtFFTFilter::ExecuteData() ERROR: Number of tuples in the input dataset is <= 1." <<std::endl;
    debug3 << "avtFFTFilter::ExecuteData() ERROR: Not executing FFT. Returning input dataset." << std::endl;
    return in_dr;
  }

  // Get the Y Axis values
  std::vector<double> yValues;
  yValues.reserve(nPts);
  for (unsigned int i = 0; i < nPts; ++i) {
    yValues.push_back( *(scalars->GetTuple(i)) );
  }  

  // Get the X Axis values
  vtkDataArray *X = rGrid->GetXCoordinates();

  // Compute dt. 
  double dt = X->GetTuple1(1) - X->GetTuple1(0);

#if 0
  //Dummy code that creates a straight line
  vtkFloatArray* vals = vtkFloatArray::New();
  vals->SetNumberOfComponents(1);
  vals->SetNumberOfTuples(nPts);
  vtkRectilinearGrid* outGrid = vtkVisItUtility::Create1DRGrid(nPts, VTK_FLOAT);
  outGrid->GetPointData()->SetScalars(vals);
  vtkFloatArray* outGrid_X = vtkFloatArray::SafeDownCast(outGrid->GetXCoordinates());
  for (int i = 0; i < nPts; i++) {
    double xVal = i;
    double yVal = i;
    outGrid_X->SetValue(i, xVal);
    vals->SetValue(i, yVal);
  }
  return outGrid;
#else
  //The actual algorithm

  // Compute FFT, passing in a pointer to the output data (outrGrid)
  vtkRectilinearGrid* outGrid = ComputeFFT(&yValues, nPts, dt);

  // Returning output dataset.
  debug5 <<"avtFFTFilter::ExecuteData() - Returning" <<std::endl;

  avtDataRepresentation *out_dr = new avtDataRepresentation(outGrid,
      in_dr->GetDomain(), in_dr->GetLabel());

  outGrid->Delete();

  return out_dr;
#endif
}

// ****************************************************************************
// Method: avtFFTFilter::UpdateDataObjectInfo
//
// Purpose:
//   Update the data object information.
//
// Note:       Work partially supported by DOE Grant SC0007548.
//
// Programmer: Brad Whitlock
// Creation:   Tue Mar 18 10:53:05 PDT 2014
//
// Modifications:
//
// ****************************************************************************

void
avtFFTFilter::UpdateDataObjectInfo(void)
{
    avtPluginDataTreeIterator::UpdateDataObjectInfo();

    GetOutput()->GetInfo().GetAttributes().AddFilterMetaData("FFT");
}
